"""
Python ↔ 油猴双向通信工具
基于localStorage + 时间戳的双向通信机制

功能特性：
1. 通过油猴扩展实现与网页的双向通信
2. 使用localStorage + 时间戳避免重复消息处理
3. 多线程架构：UI主线程 + 浏览器通信子线程
4. 自动清理浏览器进程，避免资源泄漏
"""

import asyncio  # 异步编程支持
import time     # 时间戳生成
import json     # JSON数据处理
import os       # 文件路径操作
import threading  # 多线程支持
import queue    # 线程间消息队列
import tkinter as tk  # GUI界面框架
from tkinter import ttk, scrolledtext, messagebox  # GUI组件
from datetime import datetime  # 时间格式化
import atexit  # 程序退出清理
from playwright.async_api import async_playwright  # 浏览器自动化

DEBUG_MODE = True

class CommunicationWorker(threading.Thread):
    """
    通信子线程：处理浏览器启动、轮询、数据交互
    
    职责：
    - 启动浏览器并加载油猴扩展
    - 轮询localStorage获取油猴消息
    - 向油猴发送回复消息
    - 管理浏览器生命周期
    """
    
    def __init__(self, extension_path, target_url, msg_queue):
        """
        初始化通信子线程
        
        参数：
        - extension_path: 油猴扩展路径
        - target_url: 目标网页URL
        - msg_queue: 消息队列，用于与UI主线程通信
        """
        super().__init__(daemon=True)  # 守护线程：主线程退出时自动终止
        self.extension_path = extension_path  # 油猴扩展路径
        self.target_url = target_url  # 目标网页URL
        self.msg_queue = msg_queue  # 用于向UI主线程传递日志/状态
        self.is_running = False  # 控制轮询启停
        self.browser = None  # 浏览器实例
        self.last_tamper_timestamp = 0  # 防重复处理数据的时间戳

    def _send_log(self, log_type, content):
        """
        将日志通过队列发送给UI主线程
        
        参数：
        - log_type: 日志类型（"系统"、"错误"、"油猴→Python"、"Python→油猴"）
        - content: 日志内容
        """
        now = datetime.now().strftime("%H:%M:%S")
        self.msg_queue.put(("log", log_type, now, content))

    def _send_status(self, is_running):
        """
        将通信状态通过队列发送给UI主线程
        
        参数：
        - is_running: 通信状态（True=运行中，False=已停止）
        """
        self.msg_queue.put(("status", is_running))

    def _cleanup_browser(self):
        """清理无头浏览器进程"""
        if self.browser:
            try:
                if self.browser.is_connected():
                    self._send_log("系统", "正在关闭无头浏览器...")
                    self.browser.close()
                    self._send_log("系统", "无头浏览器已关闭，无进程残留")
                else:
                    self._send_log("系统", "浏览器连接已断开，无需清理")
            except RuntimeError as e:
                if "Event loop is closed" in str(e):
                    self._send_log("系统", "事件循环已关闭，浏览器进程已自动清理")
                else:
                    self._send_log("错误", f"清理浏览器时出错：{str(e)}")
            except Exception as e:
                self._send_log("错误", f"清理浏览器时出错：{str(e)}")
            finally:
                self.browser = None  # 确保引用被清除

    async def _async_get_localstorage(self, page, key):
        """
        异步读取localStorage数据并解析
        
        参数：
        - page: 浏览器页面对象
        - key: localStorage键名
        
        返回：
        - 解析后的JSON对象，如果不存在则返回None
        """
        try:
            # 检查页面是否有效
            if page.is_closed():
                raise Exception("页面已关闭")
            
            # 通过Playwright执行JavaScript读取localStorage
            data_str = await page.evaluate(f'() => localStorage.getItem("{key}")')
            return json.loads(data_str) if data_str else None
        except Exception as e:
            # 如果是页面关闭错误或执行上下文被销毁，重新抛出以便上层处理
            if "closed" in str(e).lower() or "页面已关闭" in str(e) or "execution context was destroyed" in str(e).lower():
                raise e
            self._send_log("错误", f"读取{key}失败：{str(e)}")
            return None

    async def _async_set_localstorage(self, page, key, content):
        """
        异步向localStorage写入带时间戳的数据
        
        参数：
        - page: 浏览器页面对象
        - key: localStorage键名
        - content: 要存储的数据内容
        
        数据格式：{"data": content, "timestamp": 毫秒级时间戳}
        """
        try:
            # 检查页面是否有效
            if page.is_closed():
                raise Exception("页面已关闭")
            
            now_ts = int(time.time() * 1000)  # 毫秒级时间戳（与油猴一致）
            data = json.dumps({"data": content, "timestamp": now_ts})
            # 直接存储JSON字符串，避免双重JSON.stringify
            await page.evaluate(f'() => localStorage.setItem("{key}", `{data}`)')
            # 调试：验证数据是否写入成功
            stored_data = await page.evaluate(f'() => localStorage.getItem("{key}")')
            if stored_data:
                self._send_log("调试", f"成功写入localStorage[{key}]：{stored_data}")
            else:
                self._send_log("错误", f"localStorage写入验证失败[{key}]")
        except Exception as e:
            # 如果是页面关闭错误，重新抛出以便上层处理
            if "closed" in str(e).lower() or "页面已关闭" in str(e):
                raise e
            self._send_log("错误", f"写入{key}失败：{str(e)}")

    async def _async_clear_localstorage(self):
        """
        异步清除localStorage中的通信数据
        
        功能：
        - 清除两个通信相关的localStorage项：tampermonkey_to_py 和 py_to_tampermonkey
        - 验证清除操作是否成功
        - 处理页面关闭等异常情况
        
        异常处理：
        - 页面关闭异常：重新抛出供上层处理
        - 其他异常：记录错误日志
        """
        try:
            # 检查页面是否有效，避免在已关闭页面上操作
            if self.page.is_closed():
                raise Exception("页面已关闭")
            
            # 清除两个通信相关的localStorage项
            await self.page.evaluate('() => localStorage.removeItem("tampermonkey_to_py")')
            await self.page.evaluate('() => localStorage.removeItem("py_to_tampermonkey")')
            
            # 验证清除是否成功：检查两个键是否都不存在
            tamper_data = await self.page.evaluate('() => localStorage.getItem("tampermonkey_to_py")')
            py_data = await self.page.evaluate('() => localStorage.getItem("py_to_tampermonkey")')
            
            if not tamper_data and not py_data:
                self._send_log("调试", "localStorage清除验证成功")
            else:
                self._send_log("警告", "localStorage清除可能不完整")
                
        except Exception as e:
            # 如果是页面关闭错误，重新抛出以便上层处理
            if "closed" in str(e).lower() or "页面已关闭" in str(e):
                raise e
            self._send_log("错误", f"清除localStorage失败：{str(e)}")

    def send_to_tampermonkey(self, content):
        """
        从UI手动发送消息给油猴
        
        功能：
        - 提供主线程调用接口，允许从UI发送消息
        - 使用线程安全的方式调用异步方法
        - 处理跨线程异步操作
        
        技术实现：
        - 使用asyncio.run_coroutine_threadsafe实现线程安全
        - 获取子线程的事件循环进行异步调用
        - 设置5秒超时防止阻塞
        
        参数：
        - content: 要发送给油猴的消息内容
        """
        if not (self.is_running and hasattr(self, 'page') and self.page):
            self._send_log("错误", "通信未启动，无法发送消息")
            return
        try:
            # 使用线程安全的方式调用异步方法
            # 获取子线程的事件循环
            if hasattr(self, '_thread_loop') and self._thread_loop:
                loop = self._thread_loop
            else:
                # 如果没有保存子线程的事件循环，使用默认方法
                loop = asyncio.get_event_loop()
            
            # 使用run_coroutine_threadsafe并等待结果
            future = asyncio.run_coroutine_threadsafe(
                self._async_set_localstorage(self.page, "py_to_tampermonkey", content),
                loop
            )
            # 等待异步操作完成，设置5秒超时
            future.result(timeout=5)
            self._send_log("Python→油猴", content)
        except Exception as e:
            self._send_log("错误", f"发送消息失败：{str(e)}")

    def stop(self):
        """停止通信（外部调用）"""
        self.is_running = False
        self._send_log("系统", "正在停止通信...")

    def run(self):
        """
        子线程核心逻辑：启动浏览器+轮询
        
        功能：
        - 在子线程中运行浏览器和通信逻辑
        - 创建独立的事件循环用于异步操作
        - 处理浏览器启动、页面导航、轮询等核心功能
        
        执行流程：
        1. 检查油猴扩展路径
        2. 创建新的事件循环
        3. 异步启动浏览器和通信
        4. 异常处理和资源清理
        
        注意：此方法在子线程中运行，不能直接操作UI
        """
        self.is_running = True
        self._send_log("系统", "通信子线程已启动，初始化浏览器中...")

        try:
            # 检查油猴扩展路径是否存在
            if not os.path.exists(self.extension_path):
                self._send_log("错误", f"油猴扩展路径不存在：{self.extension_path}")
                self.is_running = False
                self._send_status(False)
                return

            # 创建新的事件循环用于异步操作（子线程需要独立的事件循环）
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            
            # 保存事件循环引用，供主线程使用（用于跨线程异步调用）
            self._thread_loop = loop
            
            # 异步启动浏览器和通信逻辑
            loop.run_until_complete(self._async_run())

        except Exception as e:
            self._send_log("错误", f"通信异常：{str(e)}")
        finally:
            # 确保资源清理
            self._cleanup_browser()
            self.is_running = False
            self._send_status(False)  # 通知UI：通信已停止
            self._send_log("系统", "通信子线程已终止")

    async def _async_run(self):
        """
        异步运行浏览器和轮询逻辑
        
        功能：
        - 使用Playwright启动浏览器并加载油猴扩展
        - 导航到目标网页并等待油猴脚本加载
        - 清除localStorage历史数据
        - 启动轮询机制监听油猴消息
        
        浏览器配置：
        - 使用持久化上下文避免无痕模式
        - 加载油猴扩展并禁用其他扩展
        - 禁用浏览器欢迎页面和更新检查
        
        轮询机制：
        - 每秒检查一次localStorage中的油猴消息
        - 通过时间戳比较避免重复处理
        - 自动回复油猴消息（可替换为AI接口）
        """
        async with async_playwright() as p:
            # 设置用户数据目录，使用持久化上下文避免无痕模式
            user_data_dir = os.path.expanduser(
                r"~\AppData\Local\Tencent\QQBrowser\User Data\Default"
            )
            
            # 使用launchPersistentContext替代launch方法，保持扩展状态
            context = await p.chromium.launch_persistent_context(
                user_data_dir=user_data_dir,
                headless=False if DEBUG_MODE else True,  # 测试时设为False，可以看到浏览器界面
                args=[
                    # 加载油猴扩展并禁用其他扩展
                    f"--load-extension={self.extension_path}",
                    f"--disable-extensions-except={self.extension_path}",
                    "--disable-features=WelcomePage",  # 禁用浏览器欢迎页面
                    "--disable-component-update",  # 禁用组件更新检查
                    "--no-first-run",  # 跳过首次运行向导
                    "--no-default-browser-check",  # 禁用默认浏览器检查
                ]
            )
            # 保存浏览器和页面引用
            self.browser = context.browser
            self.page = context.pages[0] if context.pages else await context.new_page()
            
            # 导航到目标网页并等待油猴脚本加载
            await self.page.goto(self.target_url)
            await self.page.wait_for_timeout(3000)  # 等待油猴脚本加载

            self._send_log("系统", f"已打开目标网页：{self.target_url}")
            
            # 清除localStorage历史数据，避免读取上次会话的消息
            await self._async_clear_localstorage()
            self._send_log("系统", "已清除localStorage历史数据")
            
            self._send_log("系统", "通信就绪，等待油猴消息...")
            self._send_status(True)  # 通知UI：通信已启动

            # 轮询逻辑（仅在运行状态下执行）
            while self.is_running:
                try:
                    # 检查页面是否仍然有效
                    if self.page.is_closed():
                        self._send_log("错误", "浏览器页面已关闭，停止轮询")
                        break
                    
                    # 读取油猴发送的数据
                    tamper_data = await self._async_get_localstorage(self.page, "tampermonkey_to_py")
                    if tamper_data and tamper_data["timestamp"] > self.last_tamper_timestamp:
                        # 接收油猴消息
                        content = tamper_data["data"]
                        self._send_log("油猴→Python", content)
                        # 更新最后处理的时间戳
                        self.last_tamper_timestamp = tamper_data["timestamp"]
                except Exception as e:
                    # 捕获页面操作异常
                    if "closed" in str(e).lower() or "execution context was destroyed" in str(e).lower():
                        self._send_log("错误", "浏览器页面被刷新或关闭，尝试重新连接...")
                        
                        # 尝试重新连接：刷新页面并重新初始化
                        try:
                            # 重新导航到目标网页
                            await self.page.goto(self.target_url)
                            await self.page.wait_for_timeout(3000)  # 等待油猴脚本重新加载
                            
                            # 重新清除localStorage历史数据
                            await self._async_clear_localstorage()
                            
                            self._send_log("系统", "页面重新连接成功，继续轮询")
                            continue  # 继续轮询
                        except Exception as reconnect_error:
                            self._send_log("错误", f"重新连接失败：{str(reconnect_error)}")
                            break
                    else:
                        self._send_log("错误", f"轮询异常：{str(e)}")
                
                await asyncio.sleep(1)  # 异步轮询间隔（1秒）


class TkCommunicationUI:
    """Tkinter UI主窗口"""
    def __init__(self, root):
        self.root = root
        self.root.title("Python ↔ 油猴 双向通信工具")
        self.root.geometry("800x600")  # 窗口大小
        self.root.resizable(True, True)  # 允许拉伸

        # 配置：替换为你的油猴扩展路径和目标网页
        self.extension_path = os.path.expanduser(
            r"~\AppData\Local\Tencent\QQBrowser\User Data\Default\Extensions\dhdgffkkebhmkfjojejmpbldmpobfkfo\4.18.1_0"
        )
        # 使用更通用的URL确保油猴脚本能触发
        self.target_url = "https://chat.deepseek.com"

        # 通信相关变量
        self.msg_queue = queue.Queue()  # 用于接收子线程消息
        self.worker = None  # 通信子线程实例
        self.is_communicating = False  # 当前通信状态

        # 构建UI界面
        self._setup_ui()

        # 启动UI消息监听（处理子线程传来的日志/状态）
        self._listen_queue()
        
        # 注册程序退出时的清理函数（双重保障）
        atexit.register(self._atexit_cleanup)

    def _setup_ui(self):
        """构建UI布局（网格布局：3行3列）"""
        # 1. 日志显示区（占满顶部2行）
        self.log_label = ttk.Label(self.root, text="通信日志")
        self.log_label.grid(row=0, column=0, padx=10, pady=5, sticky="w")

        self.log_text = scrolledtext.ScrolledText(
            self.root, wrap=tk.WORD, font=("Consolas", 10)
        )
        self.log_text.grid(
            row=1, column=0, columnspan=3, padx=10, pady=0,
            sticky="nsew"
        )
        self.log_text.config(state=tk.DISABLED)  # 日志区只读

        # 2. 消息发送区（第3行前2列）
        self.msg_label = ttk.Label(self.root, text="发送消息给油猴：")
        self.msg_label.grid(row=2, column=0, padx=10, pady=10, sticky="w")

        self.msg_entry = ttk.Entry(self.root, font=("Consolas", 10))
        self.msg_entry.grid(
            row=2, column=1, padx=10, pady=10, sticky="ew"
        )
        self.msg_entry.config(state=tk.DISABLED)  # 初始禁用
        # 绑定回车键发送消息
        self.msg_entry.bind("<Return>", lambda event: self._send_msg())

        self.send_btn = ttk.Button(
            self.root, text="发送", command=self._send_msg
        )
        self.send_btn.grid(row=2, column=2, padx=10, pady=10, sticky="w")
        self.send_btn.config(state=tk.DISABLED)  # 初始禁用

        # 3. 控制按钮区（第4行）
        self.control_frame = ttk.Frame(self.root)
        self.control_frame.grid(
            row=3, column=0, columnspan=3, padx=10, pady=5
        )

        self.start_stop_btn = ttk.Button(
            self.control_frame, text="启动通信", command=self._toggle_communication
        )
        self.start_stop_btn.grid(row=0, column=0, padx=20, pady=5)

        self.clear_log_btn = ttk.Button(
            self.control_frame, text="清空日志", command=self._clear_log
        )
        self.clear_log_btn.grid(row=0, column=1, padx=20, pady=5)

        # 窗口拉伸配置（日志区随窗口放大）
        self.root.grid_rowconfigure(1, weight=1)
        self.root.grid_columnconfigure(1, weight=1)

    def _append_log(self, log_type, time_str, content):
        """向日志区追加内容（主线程调用）"""
        print(f"_append_log: {log_type}, {time_str}, {content}")
        # 不同类型日志用不同颜色区分（Tkinter通过标签实现）
        color_map = {
            "油猴→Python": "#0066CC",
            "Python→油猴": "#B22222",
            "系统": "#00868B",
            "错误": "#CC0000"
        }
        color = color_map.get(log_type, "#000000")

        # 过滤掉Tkinter不支持的Unicode字符（超出U+0000-U+FFFF范围）
        def filter_unsupported_chars(text):
            """过滤掉Tkinter不支持的Unicode字符"""
            return ''.join(
                char for char in text 
                if ord(char) <= 0xFFFF  # 只保留U+0000到U+FFFF范围的字符
            )
        
        # 将content中的换行符替换为|，简化日志显示
        simplified_content = content.replace('\n', ' | ')
        filtered_content = filter_unsupported_chars(simplified_content)
        
        # 启用日志区→插入内容→禁用只读
        self.log_text.config(state=tk.NORMAL)
        
        # 创建唯一的标签名避免冲突
        tag_name = f"{log_type}_{time_str.replace(':', '')}"
        
        # 配置标签颜色
        self.log_text.tag_configure(tag_name, foreground=color)
        
        # 插入单行日志（包含时间戳和类型），并直接应用标签
        log_line = f"[{time_str}] [{log_type}] {filtered_content}\n"
        
        # 使用insert方法直接应用标签
        self.log_text.insert(tk.END, log_line, tag_name)
        
        # 禁用日志区编辑
        self.log_text.config(state=tk.DISABLED)

        # 自动滚动到最后一行
        self.log_text.see(tk.END)

    def _clear_log(self):
        """清空日志区"""
        self.log_text.config(state=tk.NORMAL)
        self.log_text.delete(1.0, tk.END)
        self.log_text.config(state=tk.DISABLED)

    def _send_msg(self):
        """手动发送消息给油猴"""
        content = self.msg_entry.get().strip()
        if not content:
            messagebox.showwarning("警告", "消息不能为空！")
            return
        # 调用子线程的发送方法
        self.worker.send_to_tampermonkey(content)
        self.msg_entry.delete(0, tk.END)  # 清空输入框

    def _toggle_communication(self):
        """启动/停止通信"""
        if not self.is_communicating:
            # 启动通信：创建子线程并启动
            self.worker = CommunicationWorker(
                extension_path=self.extension_path,
                target_url=self.target_url,
                msg_queue=self.msg_queue
            )
            self.worker.start()
            self.start_stop_btn.config(text="停止通信")
        else:
            # 停止通信：调用子线程的停止方法
            self.worker.stop()
            self.start_stop_btn.config(text="启动通信", state=tk.DISABLED)  # 禁用按钮防止重复点击

    def _update_status(self, is_running):
        """更新UI状态（根据子线程传来的状态）"""
        self.is_communicating = is_running
        if is_running:
            # 通信启动：启用发送按钮和输入框
            self.start_stop_btn.config(text="停止通信", state=tk.NORMAL)
            self.msg_entry.config(state=tk.NORMAL)
            self.send_btn.config(state=tk.NORMAL)
        else:
            # 通信停止：禁用发送按钮和输入框
            self.start_stop_btn.config(text="启动通信", state=tk.NORMAL)
            self.msg_entry.config(state=tk.DISABLED)
            self.send_btn.config(state=tk.DISABLED)

    def _listen_queue(self):
        """监听消息队列（处理子线程传来的日志和状态）"""
        try:
            # 非阻塞读取队列（超时0.1秒，避免UI卡顿）
            msg_type, *args = self.msg_queue.get(block=False, timeout=0.1)
            if msg_type == "log":
                # 处理日志消息：args = (log_type, time_str, content)
                self._append_log(*args)
            elif msg_type == "status":
                # 处理状态消息：args = (is_running,)
                self._update_status(*args)
        except queue.Empty:
            pass  # 队列空时无操作
        finally:
            # 每隔100ms监听一次队列（UI主线程循环）
            self.root.after(100, self._listen_queue)

    def _atexit_cleanup(self):
        """程序退出时的清理函数（atexit注册）"""
        print("atexit: 正在执行程序退出清理...")
        if self.worker and self.worker.is_running:
            print("atexit: 检测到运行中的线程，正在停止...")
            # 先停止线程
            self.worker.stop()
            
            # 等待线程结束，最多等待2秒
            self.worker.join(timeout=2.0)
            
            # 如果线程仍然存活，强制清理浏览器
            if self.worker.is_alive():
                print("atexit: 线程未正常结束，强制清理浏览器进程")
                if hasattr(self.worker, 'browser') and self.worker.browser:
                    try:
                        self.worker.browser.close()
                        print("atexit: 浏览器进程已强制关闭")
                    except RuntimeError as e:
                        if "Event loop is closed" in str(e):
                            print("atexit: 事件循环已关闭，浏览器进程已自动清理")
                        else:
                            print(f"atexit: 强制关闭浏览器失败: {e}")
                    except Exception as e:
                        print(f"atexit: 强制关闭浏览器失败: {e}")
            else:
                print("atexit: 线程已正常结束")
        else:
            print("atexit: 没有运行中的线程，无需清理")

    def on_close(self):
        """关闭窗口时的清理操作"""
        if self.worker and self.worker.is_running:
            # 先停止线程
            self.worker.stop()
            
            # 等待线程结束，最多等待3秒
            self.worker.join(timeout=3.0)
            
            # 如果线程仍然存活，强制清理浏览器
            if self.worker.is_alive():
                self._append_log("系统", datetime.now().strftime("%H:%M:%S"), "线程未正常结束，强制清理浏览器进程")
                if hasattr(self.worker, 'browser') and self.worker.browser:
                    try:
                        self.worker.browser.close()
                    except:
                        pass
        
        # 注销atexit函数，避免重复执行
        atexit.unregister(self._atexit_cleanup)
        self.root.destroy()


if __name__ == "__main__":
    # 检查Playwright依赖（首次运行需执行：playwright install chrome）
    try:
        import asyncio
        from playwright.async_api import async_playwright
        
        async def check_playwright():
            async with async_playwright() as p:
                pass  # 验证Playwright是否可用
        
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        loop.run_until_complete(check_playwright())
    except Exception:
        print("提示：请先执行 `playwright install chrome` 安装Chromium内核！")
        exit(1)

    # 启动UI
    root = tk.Tk()
    app = TkCommunicationUI(root)
    root.protocol("WM_DELETE_WINDOW", app.on_close)  # 关闭窗口时触发清理
    root.mainloop()